'''
Author: Angyi
Date: 2020-11-18 09:49:42
LastEditTime: 2020-11-20 10:37:08
LastEditors: Please set LastEditors
Description: Eddy 预测
FilePath: /Eddy/predict.py
'''
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
from keras.models import load_model
from train import dice_coef_loss,mean_dice_coef,weighted_mean_dice_coef


def predict(days):
    width = 168
    height = 168
    model_path = 'model/eddy.h5'
   
    model = load_model(model_path,custom_objects={'dice_coef_loss':dice_coef_loss,'mean_dice_coef':mean_dice_coef,'weighted_mean_dice_coef':weighted_mean_dice_coef})

    SSH_test = np.expand_dims(np.load('model_data/test/ssha_test.npy'),3) #you can use all the region, in here I selected a box of 168*168 size
    SSH_test = np.nan_to_num(SSH_test)
    Seg_test = np.expand_dims(np.load('model_data/test/seg_test.npy'),3)

    # SSH_train = np.expand_dims(np.load('model_data/train/ssha_train.npy'),3)

    # SSH_test = np.nan_to_num(SSH_train)
    # Seg_test = np.expand_dims(np.load('model_data/train/seg_train.npy'),3)[:,:,:,:]


    randindex=days #np.random.randint(0,len(SSH_test))
    predictedSEGM=model.predict(np.reshape(SSH_test[randindex,:,:],(1,height,width,1)))
    predictedSEGMimage = np.reshape(predictedSEGM.argmax(2),(height,width))

    fig = plt.figure(figsize=(20, 10))

    plt.subplot(131)
    plt.imshow(SSH_test[randindex,:,:,0], cmap=plt.cm.jet)
    plt.colorbar(fraction=0.042, pad=0.04)
    #plt.clim(-0.25,0.25)
    plt.axis('off')
    ax = plt.gca()
    ax.invert_yaxis()
    plt.title('SSH')

    plt.subplot(132)
    plt.imshow(predictedSEGMimage,  cmap=plt.cm.jet)
    # plt.colorbar(extend='both', fraction=0.042, pad=0.04)
    #plt.clim(-0.25,0.25)
    ax = plt.gca()
    ax.invert_yaxis()
    plt.axis('off')
    plt.title('Eddynet Segmentation')

    plt.subplot(133)
    plt.imshow(Seg_test[randindex,:,:,0],  cmap=plt.cm.jet)
    # plt.colorbar(extend='both', fraction=0.042, pad=0.04)
    #plt.clim(-0.25,0.25)
    plt.axis('off')
    ax = plt.gca()
    ax.invert_yaxis()
    plt.title('ground truth Segmentation')

    # plt.savefig('result_eddy.png')
    return fig,SSH_test[randindex,:,:,0],predictedSEGMimage,Seg_test[randindex,:,:,0]
    # return predictedSEGMimage


def predict_app(pre_input):

    width = 168
    height = 168
    model_path = 'model/eddy.h5'
   
    model = load_model(model_path,custom_objects={'dice_coef_loss':dice_coef_loss,'mean_dice_coef':mean_dice_coef,'weighted_mean_dice_coef':weighted_mean_dice_coef})
    predictedSEGM=model.predict(np.reshape(pre_input,(1,height,width,1)))
    predictedSEGMimage = np.reshape(predictedSEGM.argmax(2),(height,width))
    return predictedSEGMimage